import argparse

parser = argparse.ArgumentParser()

# ========================= Train and trainer ==========================
parser.add_argument('--warmup_steps', type=int, default=-1)
parser.add_argument('--warmup_proportion', default=0.05, type=float)
parser.add_argument('--weight_decay', default=0.01, type=float)
parser.add_argument('--lr', default=8e-5, type=float, help='initial learning rate')
parser.add_argument('--batch_size', default=16, type=int)
parser.add_argument('--eval_start', default=0, type=int)
parser.add_argument('--max_epochs', default=8, type=int)
parser.add_argument('--accumulate_grad_batches', default=1, type=int)
parser.add_argument('--max_steps', default=-1, type=int)
parser.add_argument('--seed', default=12, type=int)
parser.add_argument('--precision', default=32, type=int)
parser.add_argument('--plugins', type=str, default='ddp_sharded')
parser.add_argument('--gpus', type=int, default=1)
parser.add_argument('--kfold', type=int, default=1)
parser.add_argument('--eval_step', default=1, type=int)
parser.add_argument('--device', type=str, default='cuda')


# ========================= Data ==========================
parser.add_argument('--train_file', type=str, required=True, default='./data/train_sample.csv')
parser.add_argument('--dev_file', type=str, required=False)
parser.add_argument('--predict_file', type=str, required=False)
parser.add_argument('--num_works', type=int, default=1)
parser.add_argument('--max_length', default=512, type=int)

# ========================= Generation task ==========================
parser.add_argument('--max_source_length', default=512, type=int)
parser.add_argument('--max_target_length', default=150, type=int)
parser.add_argument('--beams', default=3, type=int)
parser.add_argument('--word_drop_prob', type=float, default=0.1)
parser.add_argument('--sent_shuffle_k ', type=float, default=2)

parser.add_argument('--m_drop', default=False, action='store_true')

# ========================= Ner task ==========================
parser.add_argument('--method', type=str, default='bio', choices=['bio', 'bios', 'span', 'biaffine'])

# ========================= Cls task ==========================
parser.add_argument('--multi_label', default=False, action='store_true')

# ========================= Model ==========================
parser.add_argument('--model_path', type=str, default='./model')
parser.add_argument('--vocab_path', type=str, default='./new_vocab.txt')
parser.add_argument('--model_type', type=str, default='bert')

# ========================= Dirs ==========================
parser.add_argument('--cache_dir', type=str, default='./cache')
parser.add_argument('--log_dir', type=str, default='./log')
parser.add_argument('--save_dir', type=str, default='./save')


# ========================= Predict ==========================
parser.add_argument('--best_model', type=str, default='')